# ============================================================================
# 全局导入
# ============================================================================
import os
import cv2
import socket
import subprocess
import os.path as osp
import numpy as np
import time
import json
import argparse
import acl
import paramiko
import logging

# ============================================================================
# 模块 1 & 2: 废弃的模块 (为了文件完整性而保留)
# ============================================================================
class DevClass:
    pass

def video_to_images(vid_file, img_folder=None, return_info=False):
    pass


# ============================================================================
# 模块 3: AI 推理 - YOLOv8_ACL_Detector 类已修改
# ============================================================================
ACL_MEM_MALLOC_HUGE_FIRST = 0
ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST = 2
ACL_SUCCESS = 0

def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    shape = im.shape[:2];
    if isinstance(new_shape, int): new_shape = (new_shape, new_shape)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup: r = min(r, 1.0)
    ratio = r, r
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]
    if auto: dw, dh = np.mod(dw, stride), np.mod(dh, stride)
    dw /= 2; dh /= 2
    if shape[::-1] != new_unpad: im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
    return im, ratio, (dw, dh)

def preprocess(image_path, model_input_size=(640, 640)):
    logging.info(f"预处理图像: {image_path}，目标尺寸 {model_input_size}")
    img0 = cv2.imread(image_path)
    if img0 is None: raise FileNotFoundError(f"无法读取图像文件(文件可能已损坏): {image_path}")
    img, _, _ = letterbox(img0, new_shape=model_input_size, stride=32, auto=False)
    img = img.transpose((2, 0, 1))[::-1]
    img = np.ascontiguousarray(img, dtype=np.float32) / 255.0
    img = np.expand_dims(img, axis=0)
    logging.info("预处理完成。")
    return img

def postprocess(preds, conf_thres=0.25, iou_thres=0.45):
    preds = preds.transpose(0, 2, 1)
    boxes, scores, class_ids = [], [], []
    for det in preds[0]:
        class_score = det[4:]
        max_score_index, max_score = np.argmax(class_score), np.max(class_score)
        if max_score > conf_thres:
            scores.append(float(max_score))
            class_ids.append(max_score_index)
            cx, cy, w, h = det[:4]
            x, y = int(cx - w / 2), int(cy - h / 2)
            boxes.append([x, y, int(w), int(h)])
    indices = cv2.dnn.NMSBoxes(boxes, scores, score_threshold=conf_thres, nms_threshold=iou_thres)
    final_class_ids = []
    if len(indices) > 0:
        indices_flat = indices.flatten()
        final_class_ids = [class_ids[i] for i in indices_flat]
    return final_class_ids

class YOLOv8_ACL_Detector:
    def __init__(self, model_path):
        self.model_path = model_path
        self.model_id = None; self.model_desc = None; self.context = None; self.stream = None
        ## MODIFIED: 不再在初始化时创建和持有这些对象
        self.input_buffer, self.output_buffers = None, []
        self.input_buffer_size, self.output_sizes = 0, []

    def init(self):
        logging.info("初始化ACL检测器资源 (Context, Stream, Model)...")
        self.context, ret = acl.rt.create_context(0)
        if ret != ACL_SUCCESS: raise RuntimeError(f"create_context 失败, ret={ret}")
        self.stream, ret = acl.rt.create_stream()
        if ret != ACL_SUCCESS: raise RuntimeError(f"create_stream 失败, ret={ret}")
        logging.info(f"从文件加载模型: {self.model_path}")
        self.model_id, ret = acl.mdl.load_from_file(self.model_path)
        if ret != ACL_SUCCESS: raise RuntimeError(f"load_from_file 失败, ret={ret}")
        self.model_desc = acl.mdl.create_desc()
        ret = acl.mdl.get_desc(self.model_desc, self.model_id)
        if ret != ACL_SUCCESS: raise RuntimeError(f"get_desc 失败, ret={ret}")
        
        ## MODIFIED: _init_buffers 被重构，只获取尺寸和创建长期buffer
        self._get_buffer_sizes_and_create_long_term_buffers()
        logging.info("检测器资源初始化成功。")

    def _get_buffer_sizes_and_create_long_term_buffers(self):
        ## MODIFIED: 这个函数现在只做它名字说的事情
        self.input_buffer_size = acl.mdl.get_input_size_by_index(self.model_desc, 0)
        self.input_buffer, ret = acl.rt.malloc(self.input_buffer_size, ACL_MEM_MALLOC_HUGE_FIRST)
        if ret != ACL_SUCCESS: raise RuntimeError(f"为输入 malloc 失败, ret={ret}")

        num_outputs = acl.mdl.get_num_outputs(self.model_desc)
        for i in range(num_outputs):
            size = acl.mdl.get_output_size_by_index(self.model_desc, i)
            self.output_sizes.append(size)
            buffer, ret = acl.rt.malloc(size, ACL_MEM_MALLOC_HUGE_FIRST)
            if ret != ACL_SUCCESS: raise RuntimeError(f"为输出 {i} malloc 失败, ret={ret}")
            self.output_buffers.append(buffer)

    def run(self, input_data):
        input_dataset = None
        output_dataset = None
        try:
            # ## MODIFIED: 在每次运行时都创建和销毁输入输出数据集 ##
            # 1. 创建输入数据集
            input_dataset = acl.mdl.create_dataset()
            data_buf_in = acl.create_data_buffer(self.input_buffer, self.input_buffer_size)
            acl.mdl.add_dataset_buffer(input_dataset, data_buf_in)

            # 2. 创建输出数据集
            output_dataset = acl.mdl.create_dataset()
            for i, size in enumerate(self.output_sizes):
                data_buf_out = acl.create_data_buffer(self.output_buffers[i], size)
                acl.mdl.add_dataset_buffer(output_dataset, data_buf_out)

            # 3. 执行核心逻辑
            ptr = acl.util.bytes_to_ptr(input_data.tobytes())
            ret = acl.rt.memcpy(self.input_buffer, self.input_buffer_size, ptr, self.input_buffer_size, ACL_MEMCPY_HOST_TO_DEVICE)
            if ret != ACL_SUCCESS: raise RuntimeError(f"memcpy H2D 失败, ret={ret}")
            
            ret = acl.mdl.execute(self.model_id, input_dataset, output_dataset)
            if ret != ACL_SUCCESS: raise RuntimeError(f"execute 失败, ret={ret}")
            
            outputs = []
            for i, size in enumerate(self.output_sizes):
                host_buffer, ret = acl.rt.malloc_host(size)
                if ret != ACL_SUCCESS: raise RuntimeError(f"malloc_host 失败, ret={ret}")
                try:
                    ret = acl.rt.memcpy(host_buffer, size, self.output_buffers[i], size, ACL_MEMCPY_DEVICE_TO_HOST)
                    if ret != ACL_SUCCESS: raise RuntimeError(f"memcpy D2H 失败, ret={ret}")
                    output_bytes = acl.util.ptr_to_bytes(host_buffer, size)
                    dims_info, _ = acl.mdl.get_cur_output_dims(self.model_desc, i)
                    shape = tuple(dims_info['dims'])
                    np_array = np.frombuffer(output_bytes, dtype=np.float32).reshape(shape)
                    outputs.append(np_array)
                finally:
                    acl.rt.free_host(host_buffer)
            return outputs
        finally:
            # 4. ## MODIFIED: 确保临时的dataset对象总被销毁 ##
            if input_dataset:
                # 销毁数据集内的 databuffer
                for i in range(acl.mdl.get_dataset_num_buffers(input_dataset)):
                    data_buf_to_destroy = acl.mdl.get_dataset_buffer(input_dataset, i)
                    acl.destroy_data_buffer(data_buf_to_destroy)
                acl.mdl.destroy_dataset(input_dataset)
            
            if output_dataset:
                for i in range(acl.mdl.get_dataset_num_buffers(output_dataset)):
                    data_buf_to_destroy = acl.mdl.get_dataset_buffer(output_dataset, i)
                    acl.destroy_data_buffer(data_buf_to_destroy)
                acl.mdl.destroy_dataset(output_dataset)


    def release(self):
        logging.info("正在释放检测器资源...")
        if self.model_id: acl.mdl.unload(self.model_id); self.model_id = None
        if self.model_desc: acl.mdl.destroy_desc(self.model_desc); self.model_desc = None
        
        # ## MODIFIED: 只销毁长期持有的buffer
        if self.input_buffer: acl.rt.free(self.input_buffer); self.input_buffer = None
        for i in range(len(self.output_buffers)):
            if self.output_buffers[i]: acl.rt.free(self.output_buffers[i]); self.output_buffers[i] = None
        
        if self.stream: acl.rt.destroy_stream(self.stream); self.stream = None
        if self.context: acl.rt.destroy_context(self.context); self.context = None
        logging.info("检测器资源已释放。")
        
# ============================================================================
# 模块 4: SFTP 功能 - 智能持久化连接 (已修正)
# ============================================================================
def get_sftp_connection(sftp_client, hostname, port, username, password, timeout=15):
    """
    检查并返回一个有效的SFTP连接。如果连接丢失，则自动重新连接。
    """
    # ** BUG FIX **: 正确的检查方法是通过 sftp_client.sock.get_transport()
    if sftp_client and sftp_client.sock.get_transport().is_active():
        return sftp_client

    logging.info("SFTP 连接丢失或尚未建立，正在创建新连接...")
    if sftp_client:
        try: sftp_client.close()
        except: pass

    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        sock.connect((hostname, port))
        transport = paramiko.Transport(sock)
        transport.connect(username=username, password=password)
        new_sftp_client = paramiko.SFTPClient.from_transport(transport)
        logging.info("SFTP 新连接建立并认证成功!")
        return new_sftp_client
    except Exception as e:
        logging.error(f"建立SFTP连接失败: {e}", exc_info=True)
        return None

# ============================================================================
# 日志设置函数
# ============================================================================
def setup_logging():
    log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(module)s - %(message)s')
    file_handler = logging.FileHandler("detector.log", mode='a', encoding='utf-8')
    file_handler.setFormatter(log_formatter)
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(log_formatter)
    root_logger = logging.getLogger()
    if root_logger.hasHandlers():
        root_logger.handlers.clear()
    root_logger.setLevel(logging.INFO)
    root_logger.addHandler(file_handler)
    root_logger.addHandler(console_handler)

# ============================================================================
# 主逻辑 / Main - 最终稳定版
# ============================================================================
def main():
    setup_logging()
    RTSP_URL = "rtsp://admin:Nanjing@8625@192.168.1.64:554/media/video1"
    # ... 其他参数
    LOCAL_IMAGE_DIR = "./temp_images"; LOCAL_JSON_PATH = "./result_place.json"
    CAPTURED_IMAGE_PATH = os.path.join(LOCAL_IMAGE_DIR, "captured_frame.png")
    MODEL_PATH = "/home/HwHiAiUser/table_tennis/modelnew/yolov8n_bs1.om"; MODEL_INPUT_SIZE = (640, 640)
    SFTP_HOSTNAME = "192.168.1.107"; SFTP_PORT = 22; SFTP_USERNAME = "ywt"; SFTP_PASSWORD = "passwd"
    REMOTE_JSON_PATH = "/data/minio/data/zlmediakit/live/210235C71P3254000069/personNum/result_place.json"
    os.environ['OPENCV_FFMPEG_CAPTURE_OPTIONS'] = 'rtsp_transport;tcp|stimeout;10000000|rw_timeout;10000000'

    detector = None
    acl_initialized = False
    sftp_client = None

    try:
        logging.info("工作进程启动，开始单次检测...")
        
        logging.info("正在初始化ACL进程...")
        ret = acl.init()
        if ret != ACL_SUCCESS: raise RuntimeError(f"acl.init 失败, ret={ret}")
        acl_initialized = True
        ret = acl.rt.set_device(0)
        if ret != ACL_SUCCESS: raise RuntimeError(f"set_device 失败, ret={ret}")
        detector = YOLOv8_ACL_Detector(MODEL_PATH)
        detector.init()
        logging.info("工作进程初始化完成。")

        cap = None
        # 移除了 try...except，让顶层来捕获
        logging.info(f"正在连接到RTSP流: {RTSP_URL}")
        cap = cv2.VideoCapture(RTSP_URL, cv2.CAP_FFMPEG)
        if not cap.isOpened(): raise ConnectionError("无法打开RTSP流。")
        ret, frame = cap.read()
        if not ret: raise IOError("从RTSP流读取帧失败。")
        os.makedirs(LOCAL_IMAGE_DIR, exist_ok=True)
        cv2.imwrite(CAPTURED_IMAGE_PATH, frame)
        
        input_data = preprocess(CAPTURED_IMAGE_PATH, model_input_size=MODEL_INPUT_SIZE)
        outputs = detector.run(input_data)
        detected_cls_ids = postprocess(outputs[0])
        person_count = sum(1 for cls_id in detected_cls_ids if int(cls_id) == 0)
        logging.info(f"检测完成。图像中的人数为: {person_count}")
        result_data = {"placedata": {"HumanNumber": person_count}}
        with open(LOCAL_JSON_PATH, 'w', encoding='utf-8') as f:
            json.dump(result_data, f, indent=2)
        logging.info(f"结果已生成并保存到本地: {LOCAL_JSON_PATH}")

        logging.info("准备通过SFTP上传文件...")
        sftp_client = get_sftp_connection(
            None, SFTP_HOSTNAME, SFTP_PORT, SFTP_USERNAME, SFTP_PASSWORD
        )
        
        if sftp_client:
            logging.info(f"开始上传文件 {LOCAL_JSON_PATH} 至 {REMOTE_JSON_PATH}...")
            sftp_client.put(LOCAL_JSON_PATH, REMOTE_JSON_PATH)
            logging.info("文件上传成功。")
        else:
            raise ConnectionError("无法获取有效的SFTP连接，上传失败。")

    except Exception as e:
        logging.critical("工作进程在执行期间发生严重错误!", exc_info=True)
        exit(1) # 以错误码退出
    finally:
        if cap is not None and cap.isOpened(): cap.release()
        logging.info("开始清理工作进程的全局资源...")
        if detector: detector.release()
        if sftp_client: sftp_client.close()
        if acl_initialized: acl.finalize()
        
    logging.info("单次检测工作进程执行完毕，正常退出。")


if __name__ == '__main__':
    main()